package com.xfsg.core.rocketmq.hook;

import com.yomahub.tlog.constant.TLogConstants;
import com.yomahub.tlog.context.SpanIdGenerator;
import com.yomahub.tlog.context.TLogContext;
import com.yomahub.tlog.utils.LocalhostUtil;
import org.apache.commons.lang3.StringUtils;
import org.apache.rocketmq.client.hook.SendMessageContext;
import org.apache.rocketmq.client.hook.SendMessageHook;

import java.util.Optional;

/**
 * @author ZhangYi
 * @classname ProducerRequestIdHook
 * @description TODO
 * @date 2022/2/18 22:58
 */
public class ProducerLogTraceHook implements SendMessageHook {

    private String serviceId;

    public ProducerLogTraceHook(String serviceId) {
        this.serviceId = serviceId;
    }

    @Override
    public String hookName() {
        return "ProducerLogTraceHook";
    }

    @Override
    public void sendMessageBefore(SendMessageContext context) {
        String traceId = TLogContext.getTraceId();
        if (StringUtils.isNotBlank(traceId)) {
            String appName = Optional.ofNullable(serviceId).orElse(TLogConstants.UNKNOWN);
            context.getMessage().putUserProperty(TLogConstants.TLOG_TRACE_KEY, traceId);
            context.getMessage().putUserProperty(TLogConstants.PRE_IVK_APP_KEY, appName);
            context.getMessage().putUserProperty(TLogConstants.PRE_IVK_APP_HOST, LocalhostUtil.getHostName());
            context.getMessage().putUserProperty(TLogConstants.PRE_IP_KEY, LocalhostUtil.getHostIp());
            context.getMessage().putUserProperty(TLogConstants.TLOG_SPANID_KEY, SpanIdGenerator.generateNextSpanId());
        }
    }

    @Override
    public void sendMessageAfter(SendMessageContext context) {

    }

}
